import re
from typing import Any
from pasers.paser_utils import re_filter
from langchain_text_splitters import CharacterTextSplitter

EXTRA_COL_NAME = 'meta_index'


def table_split(split_len, data_text, head_text):
    """
    这里切的是表格的数据部分（data_text），但要保证切割后的数据+表头的长度小于split_len
    :param split_len: 最大切割长度
    :param data_text: 表格的数据部分
    :param head_text: 表头字符串
    :return:
    """
    head_len = len(head_text)
    c_splitter = CharacterTextSplitter(
        separator='\n',
        chunk_size=split_len - head_len,
        chunk_overlap=0
    )
    results = []
    pattern = re.compile('[-| \n]+')
    for sub_table in c_splitter.split_text(data_text):
        txt = re_filter(head_text + sub_table)
        if len(re.sub(pattern, '', txt).strip()) > 0:
            results.append(txt)
    return results


def prepare_excel_table(table, max_len) -> tuple[list[Any], list[Any], list[Any]]:
    """
    :param table: df类型
    :param max_len: 如果超过该值则将表格划分为多个子表
    :return:
         table_ls：子表的markdown字符串列表
         extra_data_ls：单个单元格内容超过max_len的溢出内容列表，每项为 ((row, col), content)
         pos_ls：对应每个子表在table中的位置信息，格式为 [((row_start, col_start),(row_end, col_end)), ...]
    """
    # 为避免重复出现 'Unnamed:'，直接重命名列
    table = table.rename(columns=lambda x: x.replace('Unnamed:', '') if 'Unnamed:' in x else x)
    table = table.fillna("-")
    # 先得到每个子表对应的列切分区间
    col_indices = get_column_split_indices(table, max_len)
    # 根据列区间切分表格（每个切分后的小表保留原表所有行）
    split_table_ls = vert_split(col_indices, table)

    table_ls = []
    extra_data_ls = []  # 溢出数据，同时包含位置信息
    pos_ls = []
    # 遍历每个按列切分后的小表，同时记录对应列区间信息
    for (col_start, col_end), small_table in zip(col_indices, split_table_ls):
        if small_table.empty:
            continue
        # 传入小表的同时传递 col_offset 参数（即当前子表在原始表中的最左列索引）
        sub_tables, extra_data, row_positions = chunk_excel_markdown(small_table, max_len, col_start)
        table_ls.extend(sub_tables)
        extra_data_ls.extend(extra_data)
        # 对于每个拆分得到的子表，行区间来自chunk_excel_markdown，列区间保持不变
        for (r_start, r_end) in row_positions:
            pos_ls.append(((r_start, col_start), (r_end, col_end)))
    return table_ls, extra_data_ls, pos_ls


def get_column_split_indices(table, max_len):
    """
    table:DataFrame
    max_len:int
    根据每列内容及表头长度判断需要将表格按照列拆分的区间，
    返回形如 [(s0,e0), (s1,e1), ...] 的列表，表示原表中对应的列区间。
    """
    columns = table.columns.tolist()
    indices = []
    s = 0
    e = -1
    # markdown形式固定附加字符数量
    col_len = 4
    # 每行均有固定字符附加（如 | 与换行符），初始均为2
    row_len = [2 for _ in range(table.shape[0])]

    # 如果需要额外添加索引列，则计算额外长度（EXTRA_COL_NAME需要提前定义）
    extraColLen = len(EXTRA_COL_NAME) + 7
    extraRowLen = len(str(table.shape[0])) + 1

    for i, col in enumerate(columns):
        cur_col_len = len(str(col).replace('\n', '').strip()) + 7

        cur_row_len = []
        for cell in table.iloc[:, i]:
            cell_len = len(str(cell).replace('\n', '').strip()) + 1
            cur_row_len.append(cell_len)

        sum_row_len = list(map(lambda x, y: x + y, cur_row_len, row_len))
        if col_len + cur_col_len + max(sum_row_len) + extraColLen + extraRowLen < max_len:
            e = i
            col_len += cur_col_len
            row_len = sum_row_len
        elif e >= s:
            indices.append((s, e))
            col_len = cur_col_len + 4
            row_len = [2 + item for item in cur_row_len]
            s = i
            e = i
    indices.append((s, len(columns) - 1))
    return indices


def vert_split(indices, table):
    """
    根据给定的列区间拆分原表，每个子表保留原表所有行。
    若切分后的子表需要后续行合并，插入附加的索引列（EXTRA_COL_NAME）。
    """
    # 如果无需拆分，直接返回整体表格；可根据实际需求选择是否添加索引列
    if len(indices) == 1:
        return [table]
    split_table_ls = []
    for start, end in indices:
        split_df = table.iloc[:, start:end + 1].copy()
        # 添加索引列便于记录原始行号，名称 EXTRA_COL_NAME 需提前定义
        index_col = list(table.index)
        split_df.insert(0, EXTRA_COL_NAME, index_col)
        split_table_ls.append(split_df)
    return split_table_ls


def chunk_excel_markdown(table, max_len, col_offset):
    """
    将传入的表格转换为 markdown 格式字符串，若内容超长，则按行拆分为多个 markdown 字符串。
    同时返回：
        sub_tables: markdown 格式字符串列表
        extra_data: 溢出内容列表，每个元素为 ((row, col), content)
        row_positions: 每个子表对应的 (row_start, row_end) 列表
    参数 col_offset 表示当前子表对应原始表中最左列的索引，
    方便计算单元格在原表中的原始列位置（注意：若表中插入了索引列，则需要跳过该列）。
    """
    table_data = []  # 存储每一行生成的 markdown 行
    extra_data = []  # 存储溢出数据及其位置信息
    row_positions_all = []  # 记录每一行对应原始行号
    columns = table.columns.tolist()

    def traverse_header(columns):
        row_txt = ['|']
        for content in columns:
            cell_txt = str(content).replace('\n', '').strip()
            row_txt.append(cell_txt)
            row_txt.append('|')
        return ''.join(row_txt).strip()

    header_text = traverse_header(columns)
    head_split = '|' + '-|' * len(columns)
    head_text = header_text + '\n' + head_split + '\n'

    # 判断是否存在插入的索引列
    has_index_col = (columns[0] == EXTRA_COL_NAME)

    for i, (row_index, row) in enumerate(table.iterrows()):
        row_text = ['|']
        for cell_id, cell in enumerate(row):
            cell_text = str(cell).replace('\n', '').strip()
            # 判断是否为溢出数据，且跳过索引列（如果存在）
            if columns[cell_id] == EXTRA_COL_NAME:
                # 对索引列直接输出，不判断是否溢出
                row_text.append(cell_text)
                row_text.append('|')
                continue
            if len(cell_text) > max_len - len(head_text):
                # 记录溢出内容，同时计算原始的列索引：
                # 如存在额外的索引列，则当前单元格对应原始列号 = col_offset + (cell_id - 1)
                # 否则，original_column = col_offset + cell_id
                original_col = col_offset + (cell_id - 1) if has_index_col else col_offset + cell_id
                # 记录溢出：位置为 (row_index, original_col)
                extra_data.append(((row_index, original_col), re_filter(f'{columns[cell_id]}: {cell_text}')))
                # 在表格中用空内容表示该单元格
                row_text.append(' ')
                row_text.append('|')
            else:
                row_text.append(cell_text if cell_text != 'nan' else '无')
                row_text.append('|')
        table_data.append(''.join(row_text).strip())
        row_positions_all.append(row_index)  # 记录当前行对应的原始行号

    data_text = '\n'.join(table_data)
    full_table_text = head_text + data_text

    # 如果整个子表的 markdown 内容不超长，则整体返回，并返回完整的行区间
    if len(full_table_text) < max_len:
        row_start = min(row_positions_all) if row_positions_all else None
        row_end = max(row_positions_all) if row_positions_all else None
        return [full_table_text], extra_data, [(row_start, row_end)]

    # 如果超长，则按行拆分，每个子表均包含表头
    sub_tables = []
    row_positions = []
    current_rows = []
    current_row_indices = []
    header_len = len(head_text)
    current_len = header_len
    for idx, row_str in zip(row_positions_all, table_data):
        row_len = len(row_str) + 1  # 加换行符
        if current_len + row_len > max_len and current_rows:
            sub_table = head_text + '\n'.join(current_rows)
            sub_tables.append(sub_table)
            row_positions.append((min(current_row_indices), max(current_row_indices)))
            current_rows = [row_str]
            current_row_indices = [idx]
            current_len = header_len + row_len
        else:
            current_rows.append(row_str)
            current_row_indices.append(idx)
            current_len += row_len
    if current_rows:
        sub_table = head_text + '\n'.join(current_rows)
        sub_tables.append(sub_table)
        row_positions.append((min(current_row_indices), max(current_row_indices)))

    return sub_tables, extra_data, row_positions
